import argparse
import pathlib

import cv2
from pathlib import Path

# 这里我们取视频的第一帧来进行标注。注意⚠️不要使用截图，因为截图会使的图像大小不一致。

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # yolov5 strongsort root directory
VIDEO = ROOT / 'video'
class CaptureScreen:
    def __init__(self,frame,source):

        p = pathlib.Path(source)
        self.out_file = p.stem + "_screenshot.png"
        self.frame = frame
        if str(source).isnumeric():
            self.vidcap = cv2.VideoCapture(0)
        else:
            self.vidcap = cv2.VideoCapture(str(source))

    def run(self):
        success, image = self.vidcap.read()
        n = 1
        while n < self .frame:
            success, image = self. vidcap.read()
            n += 1
        cv2.imwrite(self.out_file, image)



def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', type=str, default=VIDEO / 'walk.mp4', help='要截图的视频')
    parser.add_argument('--frame', type=int , default=20, help='第几帧？')
    # parser.add_argument('--out', type=str, default= 'out.png', help='保存的截图名称')
    opt = parser.parse_args()
    return opt

if __name__ == '__main__':
    args = parse_opt()
    cp= CaptureScreen(args.frame,args.source)
    cp.run( )
